package com.example.spring_ai.a2a.server.service;

import com.example.spring_ai.a2a.server.bean.TaskResponse;
import com.example.spring_ai.a2a.server.bean.TaskStatus;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@Slf4j
public class TaskService {

    private final Map<String, TaskStatus> taskStore = new ConcurrentHashMap<>();

    public void processTask(String taskId, Map<String, String> params) {
        log.info("[server] Processing task {} with params {}", taskId, params);
        List<String> tasks = Arrays.asList(params.get("tasks").split(","));
        String progress =  "0/" + tasks.size();
        String finalRes = "";
        taskStore.put(taskId, new TaskStatus("working", progress, ""));
        try {
            // 模拟任务执行
            for (int i = 0; i < tasks.size(); i++) {
                progress = (i + 1) + "/" + tasks.size();
                //随机返回success或者failed
                String res = Math.random() < 0.5 ? "success" : "failed";
                finalRes = finalRes + "{" + tasks.get(i) + ":" + res + "}, ";
                taskStore.put(taskId, new TaskStatus("working", progress, finalRes));
                Thread.sleep(2000);
            }
            taskStore.put(taskId, new TaskStatus("completed", progress, finalRes));
        } catch (InterruptedException e) {
            taskStore.put(taskId, new TaskStatus("failed", progress,
                    Map.of("error", "任务被中断"))
            );
            Thread.currentThread().interrupt(); // 恢复中断状态
        } catch (Exception e) {
            // 处理其他异常
            taskStore.put(taskId, new TaskStatus("failed", progress,
                    Map.of("error", e.getMessage()))
            );
        }
        log.info("[server] Processing task {} end", taskId);
    }

    public Flux<TaskResponse> streamTaskStatus(String taskId) {
        return Flux.interval(Duration.ofSeconds(1))
                .map(tick -> {
                    TaskStatus taskStatus = taskStore.get(taskId);
                    if (taskStatus == null) {
                        return new TaskResponse(taskId, "not_found", null, null);
                    } else {
                        return new TaskResponse(
                                taskId,
                                taskStatus.getStatus(),
                                taskStatus.getProgress(),
                                taskStatus.getResult()
                        );
                    }
                })
                .takeUntil(response ->
                        "not_found".equals(response.getStatus()) ||
                                "completed".equals(response.getStatus()) ||
                                "failed".equals(response.getStatus())
                )
                .doOnNext(response ->
                        log.info("Sending SSE: {}", response))
                .doOnComplete(() ->
                        log.info("SSE stream completed for task {}", taskId))
                .onErrorResume(e -> {
                    log.error("Error in SSE stream for task {}", taskId, e);
                    return Flux.empty();
                });
    }
}
